#!/usr/bin/env python3
'''
未验证

训练记录：
完成调整，重新验证
20250213: 训练分数32.015，测试分数39.53，继续训练
20250214: Learning Rate: 0.0005

训练分数37.568，测试分数47.86，继续训练

20250217：Learning Rate: 0.0005，训练分数41.612，测试分数49.4，继续训练
20250218:Learning Rate: 0.00045000000000000004，训练分数45.113，测试分数55.633，继续训练
20250219:Learning Rate: 0.00045000000000000004，训练分数47.375，测试分数57.46，继续训练
20250220:Learning Rate: 0.00040500000000000003，训练分数47.646，测试分数59.73，继续训练
20250221:Learning Rate: 0.0003645，训练分数50.484，测试分数61.86，继续训练
20250222:Learning Rate: 0.0003645，训练分数54.076，测试分数，继续训练66.06，继续训练
20250224:Learning Rate: 0.00032805000000000003，训练分数57.274，测试分数70.96，继续训练
20250225:Learning Rate: 0.00032805000000000003，训练分数57.996，测试分数，继续训练70.96，继续训练，开放调度器一天
20250226:Learning Rate: 0.00032805000000000003，训练分数59.645，测试分数70.96，，排查代码，是否需要继续训练
20250227:Learning Rate: 0.00032805000000000003，训练分数60.814，测试分数78.63，继续训练，关闭调度器
20250228:Learning Rate: 0.000295245，训练分数61.971，测试分数79.3，继续训练
20250301:Learning Rate: 0.000295245，训练分数62.601，测试分数79.3，继续训练
20250302:Learning Rate: 0.0002657205，训练分数64.501，测试分数83，继续训练
20250303:Learning Rate: 0.00023914845，测试分数85.2，训练分数64.902，继续训练
:0250304：Learning Rate: 0.00023914845，测试分数91.63，训练分数未知，因为训练中断，继续训练
20250304:Learning Rate: 0.000215233605，测试分数91.63，训练分数 64.851，继续训练一天，查看学习率调度器
20250306：Learning Rate: 0.000215233605，测试分数91.63，训练分数未知，因为训练中断，继续训练
20250307：Learning Rate: 0.000193710244，测试分数91.63， 训练分数64.941，继续训练，继续训练一天
20250308:Learning Rate: 0.0001937102445，无进步，调整代码，重新训练
20250309:修改为多帧堆叠，重新训练，训练分数 48.927，测试分数57.9333，继续训练
200250310：Learning Rate: 0.0005，测试分数70.83分，训练分数 61.521，继续训练
20250311：Learning Rate: 0.0005，测试分数73.93，训练分数 63.320，继续训练
20250312：Learning Rate: 0.0005
20250312:测试分数73。93，训练分数 65.411，继续训练
20250313:Learning Rate: 0.0005，测试分数1044.2，训练分数 964.115，由于修改了奖励机制奖励提升较多，继续训练
20250313:Learning Rate: 0.0005，测试分数1054.16,训练分数978.413
20250314:Learning Rate: 0.0005，测试分数1057.8，训练分数972.194，继续训练
20250316:Learning Rate: 0.0005，测试分数1057.8，训练分数976.933，继续训练
20250317:继续训练
20250318：Learning Rate: 0.0005，测试分数1057，训练分数941.754，继续训练，开放学习率调度器
20250319：Learning Rate: 0.0005，测试分数1057.8，训练分数 968.337，关闭学习率调度器，继续训练
20250320：Learning Rate: 0.00045，测试分数1057.8，训练分数980.817，查看学习率调度器，继续训练
20250321:Learning Rate: 0.00045，测试分数1057.8，训练分数990.965，继续训练
20250322:Learning Rate: 0.00045
20250323:Learning Rate: 0.00045，测试分数1057.8，训练分数997.151，继续训练，开启学习率调度器
20250324:Learning Rate: 0.00045,异常中断
20250325:Learning Rate: 0.00045，测试分数_1057.8，训练分数897.316，确认是否需要继续训练
20250326:Learning Rate: 0.000405,测试分数1057.8，训练分数990.217，继续训练，确认学习率调度器
20250327：暂停训练一天，继续开放学习率调度器,
20250328：Learning Rate: 0.0003645，测试分数1057.86，训练分数1002.142，继续训练，确认学习率调度器
202050329:Learning Rate: 0.0003645，测试分数1084.26，训练分数1015.880，继续训练
20250331:Learning Rate: 0.0003645,测试分数1084.2666666666682，训练分数1022.553，继续训练
20250401:Learning Rate: 0.0003645,测试分数1084.26，训练分数1027.620，继续训练
20250402：Learning Rate: 0.0003645，测试分数1084.2666666666682，训练分数1031.299，继续训练
20250403：Learning Rate: 0.0003280,测试分数1084.2666666666682，训练分数1030.569，开启学习率调度器
20250407：Learning Rate: 0.000295245,测试分数1087.73，训练分数1034.137，继续训练
20250408：Learning Rate: 0.0002657205，测试分数1187.4333，训练分数1036.552，继续训练
20250409：Learning Rate: 0.00023914845
20250409：Learning Rate: 0.00023914845
20250410:Learning Rate: 0.000215233605
20250411:暂停训练一天，未发现新的进步，如果继续一天没有则停止训练，play模型
20250414:暂停训练一天
20250418:Learning Rate: 0.000215233605,测试分数1187.43，训练分数1034.669，没有进步，停止训练
'''
import gymnasium as gym
import ptan
import numpy as np
import argparse
from tensorboardX import SummaryWriter
import os

import torch
import torch.nn as nn
import torch.nn.utils as nn_utils
import torch.nn.functional as F
import torch.optim as optim
from collections import deque
import cv2
import ale_py


from lib import common

gym.register_envs(ale_py)
GAMMA = 0.99
LEARNING_RATE = 5e-4
ENTROPY_BETA = 0.01
BATCH_SIZE = 128
NUM_ENVS = 50

REWARD_STEPS = 4
CLIP_GRAD = 0.5

SAVE_ITERS = 100

class StackFrameWrapper(gym.Wrapper):
    def __init__(self, env, n_frames=4):
        super().__init__(env)
        self.env = env
        self.n_frames = n_frames
        self.frames = deque([], maxlen=n_frames)

        shp = env.observation_space.shape
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(shp[0]*n_frames, shp[1], shp[2]), dtype=np.float32)

        self.obs = []

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        for _ in range(self.n_frames):
            self.frames.append(obs)
        return np.concatenate(list(self.frames), axis=0), info

    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        self.frames.append(obs)
        return np.concatenate(list(self.frames), axis=0), reward, terminated, truncated, info
    

def wrap_dqn(env, stack_frames=4, episodic_life=False, reward_clipping=True):
    if episodic_life:
        # 将多条生命的游戏模拟成单条生命ActorCriticAgent
        env = ptan.common.wrappers.EpisodicLifeEnv(env)
    # 增强初始化
    env = ptan.common.wrappers.NoopResetEnv(env, noop_max=30)
    # 跳帧包装器
    # env = MaxAndSkipEnv(env, skip=4)
    env = ptan.common.wrappers.FireResetEnv(env)
    env = common.ProcessFrame84(env)
    # env = common.DenseOpticalFlow3FrameWrapper(env)
    env = ptan.common.wrappers.ImageToPyTorch(env)
    env = ptan.common.wrappers.FrameStack(env, stack_frames)
    env = common.AlienPenaltyWrapper(env)
    return env

class AtariA2C(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(AtariA2C, self).__init__()

        # Batch Normalization 移除：虽然 BatchNorm 在图像任务中很有用，但它不一定适用于强化学习的 Atari 游戏。在这些游戏中，BatchNorm 可能会扰乱学习过程，因为强化学习的输入数据是在线生成的，没有像监督学习那样的稳定分布。
        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 64, kernel_size=8, stride=4),  # kernel size and stride tuned for Atari
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=3, stride=1),
            nn.BatchNorm2d(256),
            nn.ReLU()
        )

        conv_out_size = self._get_conv_out(input_shape)
        self.policy = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions)
        )

        self.value = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        fx = x.float() / 255.0
        conv_out = self.conv(fx).view(fx.size()[0], -1)
        return self.policy(conv_out), self.value(conv_out)


def unpack_batch(batch, net, device='cpu'):
    """
    Convert batch into training tensors
    :param batch:
    :param net:
    :return: states variable, actions tensor, reference values variable
    """
    states = []
    actions = []
    rewards = []
    not_done_idx = [] # 非结束的游戏数据索引，该索引记录对应batch，states，actions，rewards
    last_states = [] # 记录采样中的执行动作后的状态，仅记录游戏非结束状态下的索引
    for idx, exp in enumerate(batch):
        states.append(np.array(exp.state, copy=False))
        actions.append(int(exp.action))
        rewards.append(exp.reward)
        if exp.last_state is not None:
            not_done_idx.append(idx)
            last_states.append(np.array(exp.last_state, copy=False))
    states_v = torch.FloatTensor(np.array(states, copy=False)).to(device)
    actions_t = torch.LongTensor(actions).to(device)
    rewards_np = np.array(rewards, dtype=np.float32)
    if not_done_idx:
        last_states_v = torch.FloatTensor(np.array(last_states, copy=False)).to(device)
        last_vals_v = net(last_states_v)[1]
        last_vals_np = last_vals_v.data.cpu().numpy()[:, 0]
        rewards_np[not_done_idx] += GAMMA ** REWARD_STEPS * last_vals_np

    ref_vals_v = torch.FloatTensor(rewards_np).to(device)
    return states_v, actions_t, ref_vals_v


@torch.no_grad()
def test_model(env, net, device, episodes=5):
    total_reward = 0.0
    for _ in range(episodes):
        obs, _ = env.reset()
        while True:
            obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
            logits_v, _ = net(obs_v)
            probs_v = F.softmax(logits_v, dim=1)
            probs = probs_v.data.cpu().numpy()
            action = np.argmax(probs)
            obs, reward, done, trunc, _ = env.step(action)
            total_reward += reward
            if done or trunc:
                break
    return total_reward / episodes


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=False, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", default="alien", required=False, help="Name of the run")
    args = parser.parse_args()
    device = common.select_device(args=args)

    save_path = os.path.join("saves", "a2c-conv-" + args.name)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    envs = [wrap_dqn(gym.make("ALE/Alien-v5", repeat_action_probability=0.0, frameskip=4, obs_type="rgb")) for _ in range(NUM_ENVS)]
    test_env = wrap_dqn(gym.make("ALE/Alien-v5", repeat_action_probability=0.0, frameskip=4, obs_type="rgb"))
    writer = SummaryWriter(comment="-alien-a2c-conv-" + args.name)

    net = AtariA2C(envs[0].observation_space.shape, envs[0].action_space.n).to(device)
    print(net)

    agent = ptan.agent.PolicyAgent(lambda x: net(x)[0], apply_softmax=True, device=device)
    exp_source = ptan.experience.ExperienceSourceFirstLast(envs, agent, gamma=GAMMA, steps_count=REWARD_STEPS)
    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50000, gamma=0.9)

    start_idx = 0
    train_count = 0
     # 增加加载模型的代码
    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)),
                             key=lambda x: int(x.split('_')[2].split('.')[0]))
        if len(checkpoints) > 0:
            checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
            net.load_state_dict(checkpoint['net'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            # 打印学习率大小
            print("Learning Rate:", scheduler.get_lr()[0])
            start_idx = checkpoint['start_idx']
            scheduler.load_state_dict(checkpoint['scheduler'])
            train_count = checkpoint["train_count"]
            print("加载模型成功")

    batch = []
    best_reward = 0
    test_best_reward = 0

    with common.RewardTracker(writer, stop_reward=10000) as tracker:
        with ptan.common.utils.TBMeanTracker(writer, batch_size=10) as tb_tracker:
            for step_idx, exp in enumerate(exp_source):
                batch.append(exp)

                new_rewards = exp_source.pop_total_rewards()
                if new_rewards:
                    if tracker.reward(new_rewards[0], step_idx + start_idx):
                        break

                if len(batch) < BATCH_SIZE:
                    continue

                states_v, actions_t, vals_ref_v = unpack_batch(batch, net, device=device)
                batch.clear()

                optimizer.zero_grad()
                logits_v, value_v = net(states_v)
                loss_value_v = F.mse_loss(value_v.squeeze(-1), vals_ref_v)

                log_prob_v = F.log_softmax(logits_v, dim=1)
                adv_v = vals_ref_v - value_v.squeeze(-1).detach()
                log_prob_actions_v = adv_v * log_prob_v[range(BATCH_SIZE), actions_t]
                loss_policy_v = -log_prob_actions_v.mean()

                prob_v = F.softmax(logits_v, dim=1)
                entropy_loss_v = ENTROPY_BETA * (prob_v * log_prob_v).sum(dim=1).mean()

                loss_policy_v.backward(retain_graph=True)
                grads = np.concatenate([p.grad.data.cpu().numpy().flatten()
                                        for p in net.parameters()
                                        if p.grad is not None])

                loss_v = entropy_loss_v + loss_value_v
                loss_v.backward()
                nn_utils.clip_grad_norm_(net.parameters(), CLIP_GRAD)
                optimizer.step()
                loss_v += loss_policy_v

                tb_tracker.track("advantage",       adv_v, step_idx + start_idx)
                tb_tracker.track("values",          value_v, step_idx + start_idx)
                tb_tracker.track("batch_rewards",   vals_ref_v, step_idx + start_idx)
                tb_tracker.track("loss_entropy",    entropy_loss_v, step_idx + start_idx)
                tb_tracker.track("loss_policy",     loss_policy_v, step_idx + start_idx)
                tb_tracker.track("loss_value",      loss_value_v, step_idx + start_idx)
                tb_tracker.track("loss_total",      loss_v, step_idx + start_idx)
                tb_tracker.track("grad_l2",         np.sqrt(np.mean(np.square(grads))), step_idx + start_idx)
                tb_tracker.track("grad_max",        np.max(np.abs(grads)), step_idx + start_idx)
                tb_tracker.track("grad_var",        np.var(grads), step_idx + start_idx)

                scheduler.step()
                train_count += 1
                if train_count % 80 == 0:
                    # Test the model
                    net.eval()
                    test_reward = test_model(test_env, net, device=device, episodes=3)
                    net.train()
                    print(f"Test reward: {test_reward:.2f}")
                    common.save_best_model(test_reward, net.state_dict(), save_path, "a2c-best", keep_best=10)


                    # Save the model if it is the best so far
                    if test_best_reward == 0 or test_reward > test_best_reward:
                        if test_best_reward != 0:
                            print(f"New best model saved with reward {test_reward:.2f}")
                        test_best_reward = test_reward

                    checkpoint = {
                        "net": net.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "start_idx": step_idx + start_idx,
                        "scheduler": scheduler.state_dict(),
                        "train_count": train_count
                    }
                    common.save_checkpoints(train_count, checkpoint, save_path, "a2c", keep_last=5)
                    print(f"save checkpoint {train_count}")
